use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::TypeEnvironment;
use super::super::TypeInference;

impl TypeInference {



    /// Infer types in an expression
    /// If expected_type is provided, use it for context-aware type inference (e.g., Result type parameters)
    pub(in crate::type_check::type_inference) fn infer_expression(&mut self, expr: &Expression, env: &mut TypeEnvironment) -> Result<Type> {
        self.infer_expression_with_expected(expr, env, None)
    }



    /// Infer types in an expression with expected type
    pub(in crate::type_check::type_inference) fn infer_expression_with_expected(
        &mut self,
        expr: &Expression,
        env: &mut TypeEnvironment,
        expected_type: Option<&Type>,
    ) -> Result<Type> {
        match expr {
            Expression::Literal(lit) => self.infer_literal(lit),
            Expression::Variable(name) => self.infer_variable_reference(name, env),
            Expression::Call(call) => {
                self.infer_call_expression_with_expected(call, env, expected_type)
            }
            Expression::MethodCall(method_call) => {
                self.infer_method_call_expression(method_call, env)
            }
            Expression::Binary(binary) => self.infer_binary_expression(binary, env),
            Expression::Unary(unary) => self.infer_unary_expression(unary, env),
            Expression::Try(try_expr) => self.infer_try_expression(try_expr, env),
            Expression::ForceUnwrap(force_expr) => {
                self.infer_force_unwrap_expression(force_expr, env)
            }
            Expression::FieldAccess(field_access) => {
                self.infer_field_access_expression(field_access, env)
            }
            Expression::IndexAccess(index_access) => {
                self.infer_index_access_expression(index_access, env)
            }
            Expression::Cast(cast) => self.infer_cast_expression(cast, env),
            Expression::Array(array) => self.infer_array_expression(array, env),
            Expression::Map(map) => self.infer_map_expression(map, env),
            Expression::Struct(struct_expr) => self.infer_struct_expression(struct_expr, env),
            Expression::If(if_expr) => self.infer_if_expression(if_expr, env),
            Expression::Match(match_expr) => self.infer_match_expression(match_expr, env),
            Expression::Lambda(lambda) => self.infer_lambda_expression(lambda, env),
            Expression::FunctionLiteral(func_lit) => {
                // Always infer from lambda if available, to ensure correct type inference
                // The function_type from parser may not have correct return type
                if let Some(lambda) = &func_lit.lambda {
                    self.infer_lambda_expression(lambda, env)
                } else if let Some(function_type) = &func_lit.function_type {
                    // Fallback to function_type if no lambda
                    Ok(function_type.clone())
                } else {
                    Err(CompilerError::type_error(
                        func_lit.location.line,
                        func_lit.location.column,
                        "FunctionLiteral must have either function_type or lambda",
                    ))
                }
            }
            Expression::SliceRange(slice_range) => {
                self.infer_slice_range_expression(slice_range, env)
            }
            Expression::Path(path_expr) => {
                self.infer_path_expression(path_expr, env, expected_type)
            }
            Expression::Rc(rc_expr) => self.infer_rc_expression(rc_expr, env),
            _ => {
                Err(CompilerError::type_error(
                    0,
                    0, // TODO: Get actual location
                    "Unsupported expression for type inference",
                ))
            }
        }
    }



    /// Infer types in a literal
    pub(in crate::type_check::type_inference) fn infer_literal(&mut self, lit: &Literal) -> Result<Type> {
        match lit {
            Literal::Integer(_) => Ok(Type::Basic(crate::ast::types::BasicType::Int)),
            Literal::Float(_) => Ok(Type::Basic(crate::ast::types::BasicType::Float64)),
            Literal::String(_) => Ok(Type::Basic(crate::ast::types::BasicType::String)),
            Literal::Char(_) => Ok(Type::Basic(crate::ast::types::BasicType::Char)),
            Literal::Boolean(_) => Ok(Type::Basic(crate::ast::types::BasicType::Bool)),
            Literal::Null => Ok(Type::Basic(crate::ast::types::BasicType::Any)),
        }
    }



    /// Infer types in a variable reference
    pub(in crate::type_check::type_inference) fn infer_variable_reference(&mut self, name: &str, env: &mut TypeEnvironment) -> Result<Type> {
        if let Some(type_) = env.lookup_variable(name) {
            Ok(type_.clone())
        } else {
            Err(CompilerError::type_error(
                0,
                0, // TODO: Get actual location
                format!("Undefined variable '{}'", name),
            ))
        }
    }

    pub(in crate::type_check::type_inference) fn infer_rc_expression(
        &mut self,
        rc_expr: &RcExpr,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        if let Some(initializer) = &rc_expr.initializer {
            let init_type = self.infer_expression(initializer, env)?;
            self.add_constraint(
                &init_type,
                rc_expr.inner_type.as_ref(),
                rc_expr.location,
            )?;
        }

        Ok(Type::Rc(RcType {
            inner_type: rc_expr.inner_type.clone(),
            location: rc_expr.location,
        }))
    }

}
